import asyncio
import sys
import os

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from tortoise import Tortoise
from config.config_dev import TORTOISE_ORM

async def reset_database():
    """
    重置数据库：删除所有表并重新创建
    """
    print("开始重置数据库...")
    try:
        # 初始化Tortoise
        await Tortoise.init(config=TORTOISE_ORM)
        print("数据库连接成功!")
        
        # 获取连接
        connection = Tortoise.get_connection("default")
        
        # 禁用外键检查
        await connection.execute_query("SET FOREIGN_KEY_CHECKS = 0")
        
        # 获取所有表名
        result = await connection.execute_query("SHOW TABLES")
        print(f"查询结果: {result}")
        
        # 根据返回结果的格式处理
        tables = []
        if len(result) > 1 and result[1]:
            # 处理查询结果，不同版本返回格式可能不同
            if isinstance(result[1], list):
                for row in result[1]:
                    if isinstance(row, dict):
                        # 从字典中获取表名
                        for key, value in row.items():
                            if key.startswith('Tables_in_'):
                                tables.append(value)
                                break
                    elif isinstance(row, (list, tuple)):
                        tables.append(row[0])
                    elif hasattr(row, '__getitem__'):
                        tables.append(row[0])
                    else:
                        tables.append(str(row))
        
        # 删除所有表
        print(f"找到 {len(tables)} 个表，正在删除...")
        for table in tables:
            try:
                await connection.execute_query(f"DROP TABLE `{table}`")
                print(f"表 '{table}' 删除成功!")
            except Exception as e:
                print(f"删除表 '{table}' 时出错: {e}")
        
        # 启用外键检查
        await connection.execute_query("SET FOREIGN_KEY_CHECKS = 1")
        print("所有表删除成功!")
        
        # 重新创建表结构
        print("正在重新创建表结构...")
        await Tortoise.generate_schemas()
        print("数据库表结构重新创建成功!")
        
    except Exception as e:
        print(f"重置数据库时出错: {e}")
        raise
    finally:
        await Tortoise.close_connections()

async def main():
    """
    主函数
    """
    print("开始重置数据库...")
    try:
        await reset_database()
        print("数据库重置完成!")
    except Exception as e:
        print(f"数据库重置失败: {e}")
        raise

if __name__ == "__main__":
    asyncio.run(main())